﻿using StableDiffusionTools.AUTOMATIC1111;
using StableDiffusionTools.Utils;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using TaggingToolkit.Events;

namespace TaggingToolkit.Framework
{
    public class ThreadedInterrogator
    {
        public event EventHandler<InterrogatedTagEventArgs> InterrogatedTag;
        protected void OnInterrogatedTag(bool success, DataTags tag, string caption, string error, int id, InterrogatedTagEventErrorSource source = InterrogatedTagEventErrorSource.None)
        {
            if (InterrogatedTag == null)
                return;

            context.Post(delegate { InterrogatedTag.Invoke(this, new InterrogatedTagEventArgs(success, tag, caption, error, id, source, scheduledItems)); }, null);
        }

        public event EventHandler InterrogatorIdle;
        protected void OnInterrogatorIdle()
        {
            if (InterrogatorIdle == null)
                return;

            context.Post(delegate { InterrogatorIdle.Invoke(this, EventArgs.Empty); }, null);
        }

        public int MaxThreads { get; set; } = 2;
        public bool Busy { get => activeThreads > 0; }

        object locker = new object();
        SynchronizationContext context;
        ConcurrentQueue<InterrogationWorkItem> queue;
        ConcurrentDictionary<DataTags, InterrogationWorkItem> schedule;
        volatile bool shutdown = false;
        volatile int activeThreads = 0;
        volatile int scheduledItems = 0;
        SemaphoreSlim apiRequestSemaphore = new SemaphoreSlim(1);

        public ThreadedInterrogator()
        {
            context = SynchronizationContext.Current;
            queue = new ConcurrentQueue<InterrogationWorkItem>();
            schedule = new ConcurrentDictionary<DataTags, InterrogationWorkItem>();
        }

        public bool ScheduleForInterrogation(List<DataTags> o, ClipType clipType)
        {
            bool result = false;
            for (int i = 0; i < o.Count; i++)
            {
                if (schedule.ContainsKey(o[i]))
                    continue;

                o[i].ScheduledForInterrogation = result = true;
                var thumb = new InterrogationWorkItem(o[i], clipType, Interlocked.Increment(ref scheduledItems));
                schedule[o[i]] = thumb;
                queue.Enqueue(thumb);
                lock (locker)
                {
                    if (activeThreads < MaxThreads && !shutdown)
                    {
                        Interlocked.Increment(ref activeThreads);
                        ThreadPool.QueueUserWorkItem(new WaitCallback(DoWork));
                    }
                }
            }
            return result;
        }

        public bool HasBeenScheduled(DataTags o)
        {
            return schedule.ContainsKey(o);
        }

        public void Stop()
        {
            lock (locker)
            {
                if (!shutdown && activeThreads > 0)
                {
                    shutdown = true;
                }
            }
        }

        private void DoWork(object obj)
        {
            while (queue.TryDequeue(out InterrogationWorkItem workItem))
            {
                if (shutdown) //bail if shutting down
                {
                    PostWorkItemStopped(workItem);
                    continue;
                }

                var base64Image = Base64Util.ConvertFileToBase64(workItem.Tag.ImagePath);

                RequestResult<string> caption;
                apiRequestSemaphore.Wait();
                if (shutdown)
                {
                    PostWorkItemStopped(workItem);
                    apiRequestSemaphore.Release();
                    continue;
                }

                try
                {
                    caption = WebUIAPI.InterrogateRequest(workItem.ClipType == ClipType.DeepBooru ? InterrogateModel.deepdanbooru : InterrogateModel.clip, base64Image).GetAwaiter().GetResult();
                }
                finally
                {
                    apiRequestSemaphore.Release();
                }

                OnInterrogatedTag(caption.Success, workItem.Tag, caption.Data, caption.Error, workItem.ID, InterrogatedTagEventErrorSource.WebRequest);
                RemoveFromSchedule(workItem);
            }


            lock (locker)
            {
                if (Interlocked.Decrement(ref activeThreads) == 0)
                {
                    shutdown = false;
                    scheduledItems = 0;
                    OnInterrogatorIdle();
                }
            }
        }

        private void RemoveFromSchedule(InterrogationWorkItem workItem)
        {
            while (!schedule.TryRemove(workItem.Tag, out var ignored)) { }
        }

        private void PostWorkItemStopped(InterrogationWorkItem workItem)
        {
            OnInterrogatedTag(false, workItem.Tag, "", "Stop was called", workItem.ID, InterrogatedTagEventErrorSource.Interrogator);
            RemoveFromSchedule(workItem);
        }
    }
}